#!/usr/bin/env python

# Constants:
PACKAGE_NAME = 'bag_player'
NODE_NAME = 'bag_player'
NEXT_SERVICE = 'bag_player_next'

TF_TYPE = 'tf/tfMessage'
IMAGE_TYPE = 'sensor_msgs/Image'
CAMERA_INFO_TYPE = 'sensor_msgs/CameraInfo'

# -------
# Imports:
# -------
# ROS
import roslib; roslib.load_manifest(PACKAGE_NAME)
import rospy, cv, rosbag
from sensor_msgs.msg import CameraInfo, Image
from tf.msg import *
from bag_player.srv import *
from roslib.msg import Clock
from argparse import ArgumentParser

# -------------
class BagPlayer:
# -------------
	# ---------------------------------
	def __init__(self, bagPath, target):
	# ---------------------------------
		# Initialize the node
		rospy.init_node(NODE_NAME)
		self.target = target

		# Load bag
		self.bag = rosbag.Bag(bagPath, 'r')

		# Start the service
		self.nextService = rospy.Service(NEXT_SERVICE, Next, self.__executeNext)

		# Get the topics 
		topicsByDataType = self.__getTopicsByDataType(self.bag)

		# Setup the publishers
		self.publishers = {}
		
		# Image Topics
		# ------------
		imageTopics = topicsByDataType[IMAGE_TYPE]
		for imageTopic in imageTopics:
			rospy.loginfo('+ Image Topic: %s' % imageTopic)
			self.publishers[imageTopic] = rospy.Publisher(imageTopic, Image)

		# Camera Info Topics
		# ------------------
		cameraInfoTopics = topicsByDataType[CAMERA_INFO_TYPE]
		for cameraInfoTopic in cameraInfoTopics:
			rospy.loginfo('+ CameraInfo Topic: %s' % cameraInfoTopic)
			self.publishers[cameraInfoTopic] = rospy.Publisher(cameraInfoTopic, CameraInfo)

		# TF Topics
		# ---------
		tfTopics = topicsByDataType[TF_TYPE]
		for tfTopic in tfTopics:
			rospy.loginfo('+ TF Topic: %s' % tfTopic)
			self.publishers[tfTopic] = rospy.Publisher(tfTopic, tfMessage)

		# Clock (this is important)
		# -------------------------
		self.clock = rospy.Publisher('/clock', Clock)

		# Combine all of the topics
		allTopics = imageTopics + cameraInfoTopics + tfTopics
		self.bagGen = self.bag.read_messages(topics = allTopics)
	
	# ----------------------------------
	def __getTopicsByDataType(self, bag): 
	# ----------------------------------
	# Copied from rxbag (just to make sure we have access to it)
		""" 
		Get all the message types in the bag and their associated topics.
		
		@param bag: bag file
		@type  bag: rosbag.Bag
		@return: mapping from message typename to list of topics
		@rtype:  dict of str to list of str
		"""
		topics_by_datatype = {}
		for c in bag._get_connections():
			topics_by_datatype.setdefault(c.datatype, []).append(c.topic)

		return topics_by_datatype

	# --------------------------
	def __executeNext(self, req):
	# --------------------------
		# Collections of published topics
		topics = []

		# Publish the next message
		while True:
			# Get the next message
			topic, msg, t = self.bagGen.next()

			# If we've seen this topic before, break
			if topic == self.target and topic in topics:
				rospy.loginfo(str(msg.header.stamp))
				break
			else: 
				topics.append(topic)

			# Publish the time 
			t = Clock()
			if topic == '/tf':
				t.clock = msg.transforms[0].header.stamp
			else:
				t.clock = msg.header.stamp
			self.clock.publish(t)


			# Lookup the publisher & republish the message
			pub = self.publishers[topic]
			pub.publish(msg)

		return NextResponse()
	
	# ------
	# Public|
	# -----------
	def run(self):
	# -----------
		rospy.loginfo('+ Running')
		rospy.spin()
		rospy.loginfo('+ Done')

# -----------------------
if __name__ == '__main__':
# -----------------------
	# Get the path to the annotation & ground truth bags from the command line arguments	
	parser = ArgumentParser(description='Compare annotations to ground truth')
	parser.add_argument('-b', '--bag', dest='bag', help='Relative path to the bag')
	parser.add_argument('-t', '--target', dest='target', help='Target topic')
	args = parser.parse_args()

	# Create the node
	node = None
	try:
		node = BagPlayer(args.bag, args.target)
		node.run()

	except IOError as e:
		print 'Error opening bag file: %s' % str(e)
